import os
import sys
import cv2
import argparse
import json
import random
import pandas as pd
import numpy as np


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', type=str, default='/tcdata')
    parser.add_argument('--output_dir', type=str, default='labels')
    parser.add_argument('--mode', type=int, default=0)
    args = parser.parse_args()
    return args


def get_direction(text):
    if text == '底部朝下':
        return 0
    elif text == '底部朝右':
        return 1 
    elif text == '底部朝上':
        return 2
    elif text == '底部朝左':
        return 3
    else:
        return -1


def decode_annotation(ann):
    if len(ann) > 1:
        default_dir = get_direction(ann[1]['option'])
    else:
        default_dir = 0
    labels = []
    for a in ann[0]:
        aj = json.loads(a['text'], encoding='utf-8')
        text = aj['text']
        direction = aj.get('direction', None)
        coord = [float(x) for x in a['coord']]
        coord = np.asarray(coord, dtype=np.float32).reshape(4, 2)
        direction = get_direction(direction)
        if direction == -1:
            direction = default_dir
        labels.append({
            'text': text,
            'points': coord,
            'direction': direction
        })

    return labels


def generate_det_annotation(csv_file, output_dir, add_prefix=True):
    data = pd.read_csv(csv_file)
    print(csv_file, 'data length:', len(data))

    csv_basename = os.path.basename(csv_file).replace('.csv', '')
    os.makedirs(output_dir, exist_ok=True)
    out_fname = os.path.join(output_dir, os.path.basename(csv_file).replace('.csv', '.txt'))
    fout = open(out_fname, 'w', encoding='utf-8')

    for i in range(len(data)):
        row = data.loc[i]
        img_name = json.loads(row[1])['tfspath'].split('/')[-1]
        ann = json.loads(row[2], encoding='utf-8')
        ann = decode_annotation(ann)

        ann_dicts = []
        for k in range(len(ann)):
            points = ann[k]['points'].round().astype('int').tolist()
            text = ann[k]['text']
            direction = ann[k]['direction']
            if len(text) == 0 or text == '*' or text == ' ':
                text = '*'
            ann_dicts.append({
                'transcription': text,
                'points': points,
                'direction': direction
            })
        if add_prefix:
            fout.write('{}/{}\t'.format(csv_basename, img_name))
        else:
            fout.write('{}\t'.format(img_name))
        fout.write(json.dumps(ann_dicts, ensure_ascii=False))
        fout.write('\n')

    fout.close()


def split_image_data(fname, output_dir, ratio=0.9):
    random.seed(100)
    with open(fname, 'r', encoding='utf8') as fh:
        lines = fh.readlines()

    num = len(lines)
    num_train = int(ratio * num)
    indices = list(range(num))
    random.shuffle(indices)
    flag = np.zeros(num, dtype=bool)
    flag[indices[:num_train]] = True

    basename = os.path.splitext(os.path.basename(fname))[0]
    os.makedirs(output_dir, exist_ok=True)
    train_fname = os.path.join(output_dir, basename + '_train.txt')
    train_file = open(train_fname, 'w', encoding='utf8')
    val_fname = os.path.join(output_dir, basename + '_val.txt')
    val_file = open(val_fname, 'w', encoding='utf8')

    for i in range(num):
        if flag[i]:
            train_file.write(lines[i])
        else:
            val_file.write(lines[i])
    train_file.close()
    val_file.close()


def main():
    args = parse_args()
    data_dir = args.data_dir
    output_dir = args.output_dir
    mode = args.mode

    csv_files = [
        os.path.join(data_dir, 'Xeon1OCR_round2_train1_20210816.csv'),
        os.path.join(data_dir, 'Xeon1OCR_round2_train2_20210816.csv'),
    ]

    for i, fname in enumerate(csv_files):
        generate_det_annotation(fname, output_dir, mode==0)
    
    # train_label_dir = os.path.join(output_dir, 'train')
    # for fname in csv_files:
    #     label_fname = os.path.join(output_dir, os.path.basename(fname).replace('.csv', '.txt'))
    #     split_image_data(label_fname, train_label_dir, ratio=0.9)


if __name__ == '__main__':
    main()
    